package edu.washington.ling.syntric.fold;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;
import java.util.Random;

import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.transform.OutputKeys;
import javax.xml.transform.Result;
import javax.xml.transform.Source;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerConfigurationException;
import javax.xml.transform.TransformerException;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;

import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.NamedNodeMap;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;
import org.xml.sax.SAXException;


public class FoldConfiguration {
	private List<String> exampleIDs;
	private List<Fold> folds;

	private FoldConfiguration() {
		exampleIDs = new LinkedList<String>();
		folds = new ArrayList<Fold>();
	}
	public static FoldConfiguration init(List<String> exampleIDs, int numFolds) {
		FoldConfiguration foldConfig = new FoldConfiguration();
		foldConfig.setExampleIDs(exampleIDs);
		foldConfig.initFolds(numFolds);
		return foldConfig;
	}
	public static FoldConfiguration init(String fileName) {
		FoldConfiguration foldConfig = new FoldConfiguration();
		foldConfig.readXMLFile(fileName);
		return foldConfig;
	}
	public void setExampleIDs(List<String> exampleIDs) {
		this.exampleIDs = exampleIDs;
	}
	public void initFolds(int numFolds) {
		folds = new ArrayList<Fold>();
		List<String> unselectedIDs = new ArrayList<String>(exampleIDs);
		List<Integer> testFoldSizes = getEvenSplit(exampleIDs.size(), numFolds);
		int foldID = 0;
		for (Integer testFoldSize: testFoldSizes) {
			folds.add(selectFoldBySampling(unselectedIDs, testFoldSize, foldID));
			foldID++;
		}
	}
	private List<Integer> getEvenSplit(int dataSetSize, int numFolds) {
		List<Integer> sizes = new LinkedList<Integer>();
		int size = dataSetSize / numFolds;
		int remainder = dataSetSize % numFolds;
		for (int i = 0; i < numFolds; i++) {
			if (remainder > 0) {
				sizes.add(size + 1);
				remainder--;
			} else {
				sizes.add(size);
			}
		}
		return sizes;
	}
	private Fold selectFoldBySampling(List<String> unselectedIDs, int testFoldSize, int foldId) {
		Random rand = new Random();
		List<String> testExampleIDs = new LinkedList<String>();
		
		for(int i = 0; i < testFoldSize; i++) {
			int index = rand.nextInt(unselectedIDs.size());
			testExampleIDs.add(unselectedIDs.get(index));
			unselectedIDs.remove(index);
		}
		Set<String> fileNameSet = new HashSet<String>(exampleIDs);
		fileNameSet.removeAll(testExampleIDs);
		List<String> trainExampleIDs = new LinkedList<String>(fileNameSet);
		return new Fold(trainExampleIDs, testExampleIDs, foldId);
	}

	public boolean checkExampleIDs(List<String> exampleIDs) {
		boolean checkSuccessful = false;
		Set<String> exampleIDsFromInput = new HashSet<String>(exampleIDs);
		Set<String> exampleIDsFromConfig = new HashSet<String>(this.exampleIDs);
		if (exampleIDsFromInput.containsAll(exampleIDsFromConfig) 
				&& exampleIDsFromConfig.containsAll(exampleIDsFromInput)) {
			checkSuccessful = true;
		}
		return checkSuccessful;
	}
	public Fold getFold(int index) {
		return folds.get(index);
	}
	public int numFolds() {
		return folds.size();
	}
	public void readXMLFile(String fileName) {
		try {
			DocumentBuilderFactory docBuilderFactory = DocumentBuilderFactory.newInstance();
            docBuilderFactory.setFeature("http://apache.org/xml/features/allow-java-encodings", 
                   true);
			DocumentBuilder docBuilder = docBuilderFactory.newDocumentBuilder();
            Document doc = docBuilder.parse(new File(fileName));

            doc.normalize();
            readDomDocument(doc);
		}
		catch(SAXException e) {
			e.printStackTrace(System.err);
		}
		catch(ParserConfigurationException e) {
			e.printStackTrace(System.err);
		}
		catch(IOException e) {
			e.printStackTrace(System.err);
		}
	}
	public void writeXMLFile(String fileName) {
		try {
			DocumentBuilderFactory docBuilderFactory = DocumentBuilderFactory.newInstance();
			DocumentBuilder docBuilder = docBuilderFactory.newDocumentBuilder();
			Document doc = docBuilder.newDocument();
			writeDomDocument(doc);
			TransformerFactory xformFactory = TransformerFactory.newInstance();
			xformFactory.setAttribute("indent-number", new Integer(2));
			Transformer xform = xformFactory.newTransformer();
			xform.setOutputProperty(OutputKeys.INDENT, "yes");
			Source source = new DOMSource(doc);
			PrintWriter writer = new PrintWriter(fileName);
			Result result = new StreamResult(writer);
			xform.transform(source, result);
			writer.flush();
			writer.close();
		} 
		catch (ParserConfigurationException e) {
			e.printStackTrace();
		}
		catch (TransformerConfigurationException e) {
			e.printStackTrace();
		}
		catch (FileNotFoundException e) {
			e.printStackTrace();
		}
		catch (TransformerException e) {
			e.printStackTrace();
		}
	}
	private void readDomDocument(Document doc) {
		folds = new LinkedList<Fold>();
		NodeList foldNodes = doc.getElementsByTagName("Fold");
		for (int i = 0; i < foldNodes.getLength(); i++) {
			Element foldNode = (Element) foldNodes.item(i);
			NamedNodeMap foldNodeAttrs = foldNode.getAttributes();
			Node idNode = foldNodeAttrs.getNamedItem("Id");
			int id = Integer.parseInt(idNode.getNodeValue());
			Element trainSetNode = (Element) foldNode.getElementsByTagName("TrainSet").item(0);
			NodeList trainExampleIDNodes = trainSetNode.getElementsByTagName("ExampleID");
			List<String> trainExampleIDs = new LinkedList<String>();
			for (int j = 0; j < trainExampleIDNodes.getLength(); j++) {
				Element trainExampleIDNode = (Element) trainExampleIDNodes.item(j);
				trainExampleIDs.add(trainExampleIDNode.getTextContent());
			}
			Element testSetNode = (Element) foldNode.getElementsByTagName("TestSet").item(0);
			NodeList testExampleIDNodes = testSetNode.getElementsByTagName("ExampleID");
			List<String> testExampleIDs = new LinkedList<String>();
			for (int j = 0; j < testExampleIDNodes.getLength(); j++) {
				Element testExampleIDNode = (Element) testExampleIDNodes.item(j);
				testExampleIDs.add(testExampleIDNode.getTextContent());
			}
			folds.add(new Fold(trainExampleIDs, testExampleIDs, id));
		}
		if (folds.size() > 0) {
			exampleIDs = folds.get(0).getExampleIDs();
		}
	}
	private void writeDomDocument(Document doc) {
		// Create root
		Element root = doc.createElement("FoldConfiguration");
		doc.appendChild(root);
		
		// Add folds
		Element foldSetNode = doc.createElement("FoldSet");
		for (Fold fold: folds) {
			Element foldNode = doc.createElement("Fold");
			foldNode.setAttribute("Id", fold.getFoldID().toString());
			Element trainSetNode = doc.createElement("TrainSet");
			for(String fileName: fold.getTrainExampleIDs()) {
				Element fileNameNode = doc.createElement("ExampleID");
				fileNameNode.setTextContent(fileName);
				trainSetNode.appendChild(fileNameNode);
			}
			foldNode.appendChild(trainSetNode);
			Element testSetNode = doc.createElement("TestSet");
			for(String fileName: fold.getTestExampleIDs()) {
				Element fileNameNode = doc.createElement("ExampleID");
				fileNameNode.setTextContent(fileName);
				testSetNode.appendChild(fileNameNode);
			}
			foldNode.appendChild(testSetNode);
			foldSetNode.appendChild(foldNode);
		}
		root.appendChild(foldSetNode);
	}
}
